package com.ark.common.mybatis.permission;

import com.ark.common.domain.BaseEntity;
import com.ark.common.domain.entity.sys.SysRole;
import com.ark.common.domain.entity.sys.SysUser;
import com.ark.common.domain.model.LoginUser;
import com.ark.common.utils.SecurityUtils;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import com.google.common.collect.Sets;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.NullValue;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.security.access.hierarchicalroles.RoleHierarchy;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;

import java.util.*;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

@Slf4j
@AllArgsConstructor
public class DataPermissionRule implements PermissionRule {

    /**
     * LoginUser 的 Context 缓存 Key
     */
    protected static final String CONTEXT_KEY = DataPermissionRule.class.getSimpleName();
    static final Expression EXPRESSION_NULL = new NullValue();
    private static final String USER_COLUMN_NAME = "user_id";
    private static final String AGENT_COLUMN_NAME = "agent_id";
    private final RoleHierarchy roleHierarchy;
    /**
     * 基于用户的表字段配置 一般情况下，每个表的部门编号字段是 dept_id，通过该配置自定义。
     * <p>
     * key：表名 value：字段名
     */
    private final Map<String, String> userColumns = new HashMap<>();
    /**
     * 所有表名，是{@link #userColumns} 的合集
     */
    private final Set<String> TABLE_NAMES = new HashSet<>();


    @Override
    public Set<String> getTableNames() {
        return TABLE_NAMES;
    }

    @Override
    public Expression getExpression(Table table) {
        LoginUser loginUser = SecurityUtils.getLoginUser();
        if (ObjectUtils.isEmpty(loginUser)){
            return EXPRESSION_NULL;
        }

        if (SysUser.isAdmin(loginUser.getUserId())){
            return EXPRESSION_NULL;
        }

        // 获取当前登录用户角色
        Set<String> roles = loginUser.getUser().getRoles().stream().map(SysRole::getRoleKey).collect(Collectors.toSet());
        Set<String> set = Sets.newHashSet(roles);
        AtomicReference<Set<String>> collect = new AtomicReference<>(Collections.emptySet());
        roles.forEach(role -> {
            Collection<? extends GrantedAuthority> collection = roleHierarchy.getReachableGrantedAuthorities(Collections.singletonList(new SimpleGrantedAuthority("ROLE_" + role)));
            collect.set(collection.stream().map(GrantedAuthority::getAuthority).map(s -> s.substring(5)).collect(Collectors.toSet()));
        });
        Sets.SetView<String> union = Sets.union(set, collect.get());
        if (union.contains("agent")){
            // 如果是代理商，则查询自己及以下创建的数据
            return getAndExpression(table, new LongValue(loginUser.getUserId()), AGENT_COLUMN_NAME);
        }
        if (union.contains("common")){
            // 如果是普通用户，则查询自己创建的数据
            return getAndExpression(table, new LongValue(loginUser.getUserId()), USER_COLUMN_NAME);
        }
        return EXPRESSION_NULL;
    }

    /**
     * 获取表达式  xxx  =  xxx 的表达式
     *
     * @return 表达式
     */
    private Expression getAndExpression(Table table, LongValue value, String columnName) {
        if (StringUtils.isBlank(columnName)) {
            columnName = userColumns.get(table.getName());
        }
        if (columnName == null) {
            log.warn("表 {} 未配置字段{}，将不会做数据权限控制", table.getName(), columnName);
            return EXPRESSION_NULL;
        }
        Column column = new Column(columnName);
        column.setTable(table);// 设置字段的表信息，避免多表联查下的错误
        return new EqualsTo(column, value);
    }

    // ==================== 添加配置 ====================
    public void addUserColumn(Class<? extends BaseEntity> entityClass) {
        addUserColumn(entityClass, USER_COLUMN_NAME);
    }

    public void addUserColumn(Class<? extends BaseEntity> entityClass, String columnName) {
        String tableName = TableInfoHelper.getTableInfo(entityClass).getTableName();
        addUserColumn(tableName, columnName);
    }

    public void addUserColumn(String tableName, String columnName) {
        userColumns.put(tableName, columnName);
        TABLE_NAMES.add(tableName);
    }
}
